#!/usr/bin/env python3
# 导入必要的库
import argparse
import os
import sys
from time import sleep
import grpc
import traceback
import threading

# 添加P4Runtime工具库路径
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 './utils/')
)
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

import struct
eth_header_format = "!6s6sH"
cpu_header_format = "!B"
eth_header_length = struct.calcsize(eth_header_format)
cpu_header_offset = eth_header_length

H2_MAC="08:00:00:00:02:02"
H22_MAC="08:00:00:00:02:22"
H2_IP="10.0.2.2"
H22_IP="10.0.2.22"
H1_MAC="08:00:00:00:01:01"
H11_MAC="08:00:00:00:01:11"
H1_IP="10.0.1.1"
H11_IP="10.0.1.11"
H3_IP="10.0.3.3"
H3_MAC="08:00:00:00:03:03"


def setEcnThresholdForSwitch(p4info_helper, egress_sw, threshold):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.check_ecn",
        match_fields={},
        action_name="MyEgress.mark_ecn",
        action_params={
            "ecn_threshold": threshold,
        })
    egress_sw.WriteTableEntry(table_entry)
    print(f"成功向交换机 {egress_sw.name} 下发ECN阈值配置：{threshold}")

def writeIngressRules(p4info_helper, sw , macdst , dport,dst_ip_addr):
    table_entry =p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr":macdst,
            "port":dport,
        })
    sw.WriteTableEntry(table_entry)

def printGrpcError(e):
    """打印gRPC通信错误的详细信息"""
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print(f"({status_code.name})", end=' ')
    traceback_obj = sys.exc_info()[2]
    print(f"[{traceback_obj.tb_frame.f_code.co_filename}:{traceback_obj.tb_lineno}]")


def get_valid_ecn_threshold():
    """获取用户输入的有效ECN阈值，输入错误时重试，直到正确为止"""
    # 循环重试：直到用户输入合法值才退出循环
    while True:
        # 1. 获取用户输入（处理输入为空的情况）
        input_str = input("请输入ECN队列长度阈值（整数，范围：0-524287）：").strip()

        # 2. 先判断输入是否为空
        if not input_str:
            print("输入不能为空，请重新输入！\n")
            continue  # 回到循环开头，让用户重新输入

        # 3. 尝试将输入转为整数，并校验范围
        try:
            threshold = int(input_str)
            # 校验阈值是否在bit<19>的合法范围（0 ~ 2^19-1 = 524287）
            if 0 <= threshold <= 524287:
                # 输入合法，返回阈值
                return threshold
            else:
                # 范围不合法，提示后重试
                print(f"输入错误：阈值必须在0-524287之间（当前输入：{threshold}），请重新输入！\n")
        except ValueError:
            # 无法转为整数（如输入字母、符号），提示后重试
            print(f"输入错误：'{input_str}' 不是有效的整数，请重新输入！\n")


def writeCloneRules(p4info_helper, sw, clone_session_id, replicas):
    SessionEntry = p4info_helper.buildCloneSessionEntry(
        clone_session_id = clone_session_id,
        replicas = replicas
    )
    sw.WritePREEntry(SessionEntry)

def fetch_responses(connection):
    try:
        for response in connection.stream_msg_resp:
            print("Received response-----")
            if response.WhichOneof("update") == "packet":
                packet = response.packet.payload
                cpu_header = packet[cpu_header_offset:cpu_header_offset+1]
                ecn=struct.unpack(cpu_header_format,cpu_header)[0]
                print("ECN value:",ecn)
                if ecn == 3:
                    print('Congestion happens!!')
                sleep(0.1)
    except grpc.RpcError as e:
        print(f"{connection.name} gRPC 错误: {e}")
    except Exception as e:
        print(f"{connection.name} 错误: {e}")

def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 创建与交换机s1的连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')

        # 创建与交换机s2的连接
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')

        # 创建与交换机s3的连接（新增）
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # 发送主控制器仲裁消息
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # 在交换机上安装P4程序
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s1")

        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s2")

        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s3")

        writeIngressRules(p4info_helper, sw=s1, macdst=H2_MAC, dport=3, dst_ip_addr=H2_IP)

        writeIngressRules(p4info_helper, sw=s1, macdst=H22_MAC, dport=3, dst_ip_addr=H22_IP)

        writeIngressRules(p4info_helper, sw=s2, macdst=H1_MAC, dport=3, dst_ip_addr=H1_IP)

        writeIngressRules(p4info_helper, sw=s2, macdst=H11_MAC, dport=3, dst_ip_addr=H11_IP)

        writeIngressRules(p4info_helper, sw=s2, macdst=H2_MAC, dport=2, dst_ip_addr=H2_IP)

        writeIngressRules(p4info_helper, sw=s2, macdst=H22_MAC, dport=1, dst_ip_addr=H22_IP)

        writeIngressRules(p4info_helper, sw=s1, macdst=H1_MAC, dport=2, dst_ip_addr=H1_IP)

        writeIngressRules(p4info_helper, sw=s1, macdst=H11_MAC, dport=1, dst_ip_addr=H11_IP)

        writeIngressRules(p4info_helper, sw=s3, macdst=H2_MAC, dport=3, dst_ip_addr=H2_IP)

        writeIngressRules(p4info_helper, sw=s3, macdst=H22_MAC, dport=3, dst_ip_addr=H22_IP)

        writeIngressRules(p4info_helper, sw=s3, macdst=H1_MAC, dport=2, dst_ip_addr=H1_IP)

        writeIngressRules(p4info_helper, sw=s3, macdst=H11_MAC, dport=2, dst_ip_addr=H11_IP)

        writeIngressRules(p4info_helper, sw=s3, macdst=H3_MAC, dport=1, dst_ip_addr=H3_IP)

        writeIngressRules(p4info_helper, sw=s2, macdst=H3_MAC, dport=4, dst_ip_addr=H3_IP)

        writeIngressRules(p4info_helper, sw=s1, macdst=H3_MAC, dport=4, dst_ip_addr=H3_IP)




        ecn_threshold = get_valid_ecn_threshold()
        setEcnThresholdForSwitch(p4info_helper,egress_sw=s1,threshold=ecn_threshold)

        setEcnThresholdForSwitch(p4info_helper,egress_sw=s2,threshold=ecn_threshold)

        setEcnThresholdForSwitch(p4info_helper,egress_sw=s3,threshold=ecn_threshold)

        writeCloneRules(p4info_helper ,s1 ,100, [{"egress_port": 252 ,"instance": 1}])

        writeCloneRules(p4info_helper ,s2 ,100, [{"egress_port": 252 ,"instance": 1}])

        writeCloneRules(p4info_helper ,s3 ,100, [{"egress_port": 252 ,"instance": 1}])

        print('\n Monitoring network congestion')

        # s1_thread = threading.Thread(
        #     target=fetch_responses,  # 目标函数
        #     args=(s1,),              # 传入参数（注意逗号，tuple格式）
        #     daemon=True              # 守护线程：主线程退出时自动退出
        # )
        # # 线程2：监听s2的消息
        # s2_thread = threading.Thread(
        #     target=fetch_responses,
        #     args=(s2,),
        #     daemon=True
        # )
        #
        # s3_thread = threading.Thread(
        #     target=fetch_responses,
        #     args=(s3,),
        #     daemon=True
        # )
        # # 启动线程
        # s1_thread.start()
        # s2_thread.start()
        # s3_thread.start()
        while True:
            sleep(1)
            print('.')
            fetch_responses(s1)
            fetch_responses(s2)
            fetch_responses(s3)
        print("监听线程已启动")

        while True:
            input("Press Enter to exit")
            break




    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    except Exception as e:
        print(f"Unexpected error: {str(e)}")
        traceback.print_exc()

    ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/ecn.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/ecn.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print(f"\np4info file not found: {args.p4info}\nHave you run 'make'?")
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print(f"\nBMv2 JSON file not found: {args.bmv2_json}\nHave you run 'make'?")
        parser.exit(1)

    main(args.p4info, args.bmv2_json)
